from typing import Optional, Dict
from langchain_chroma import Chroma
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.tools import Tool
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.embeddings.dashscope import DashScopeEmbeddings
from langchain_core.output_parsers import StrOutputParser
import os
from langchain_openai import ChatOpenAI


class KnowledgeRetrieverFactory:
    """知识检索器工厂类"""

    @staticmethod
    def create_retriever() -> 'KnowledgeRetriever':
        """创建知识检索器实例"""
        embeddings = DashScopeEmbeddings(
            model='text-embedding-v1',
            dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")
        )
        return KnowledgeRetriever(embeddings)


class KnowledgeRetriever:
    """知识检索器（严格限定三个专业领域）"""

    def __init__(self, embeddings: DashScopeEmbeddings):
        self.embeddings = embeddings
        # 严格限定三个领域
        self.available_domains = ["it_tech", "medical_care", "hr_management"]
        self.domain_mapping = {
            "it_tech": "IT技术",
            "medical_care": "医疗健康",
            "hr_management": "人力资源"
        }
        # 非专业领域提示语
        self.out_of_domain_response = "抱歉，此问题不在我的专业领域范围内（IT技术、医疗健康、人力资源），无法为您解答。"

    def get_vector_db(self, domain: str) -> Chroma:
        """连接已有向量数据库"""
        return Chroma(
            persist_directory="./chroma_db",
            embedding_function=self.embeddings,
            collection_name=f"col_{domain}"
        )

    def retrieve_knowledge(self, query: str, domain: str) -> Optional[str]:
        """知识检索方法"""
        try:
            vectordb = self.get_vector_db(domain)
            docs = vectordb.similarity_search(query, k=3)
            if docs:
                return "\n".join([
                    f"问题：{doc.metadata['question']}\n答案：{doc.metadata['answer']}"
                    for doc in docs
                ])
            return None
        except Exception as e:
            print(f"检索错误: {str(e)}")
            return None


class DomainClassifierFactory:
    """领域分类器工厂类"""

    @staticmethod
    def create_classifier(llm: ChatOpenAI, domain_mapping: Dict[str, str]) -> 'StrictDomainClassifier':
        """创建领域分类器实例"""
        return StrictDomainClassifier(llm, domain_mapping)


class StrictDomainClassifier:
    """严格领域分类器（仅限三个专业领域）"""

    def __init__(self, llm: ChatOpenAI, domain_mapping: Dict[str, str]):
        self.llm = llm
        self.domain_mapping = domain_mapping
        self.prompt = ChatPromptTemplate.from_template(
            """请严格判断以下问题是否属于IT技术、医疗健康或人力资源领域。
            如果是则返回最相关的领域名称，否则返回"other"。
            问题：{question}
            只需返回领域名称或"other"，不要解释。"""
        )

    def classify(self, question: str) -> str:
        """分类用户问题所属领域"""
        chain = self.prompt | self.llm | StrOutputParser()
        try:
            result = chain.invoke({"question": question})
            if result in self.domain_mapping.values():
                return next(k for k, v in self.domain_mapping.items() if v == result)
            return "other"
        except Exception:
            return "other"


class AgentFactory:
    """Agent工厂类"""

    @staticmethod
    def create_agent() -> 'StrictDomainAgent':
        """创建Agent实例"""
        retriever = KnowledgeRetrieverFactory.create_retriever()
        return StrictDomainAgent(retriever)


class StrictDomainAgent:
    """严格限定领域问答Agent"""

    def __init__(self, retriever: KnowledgeRetriever):
        self.llm = ChatOpenAI(
            model="qwen-max",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url='https://dashscope.aliyuncs.com/compatible-mode/v1',
            temperature=0
        )

        self.retriever = retriever
        self.classifier = DomainClassifierFactory.create_classifier(
            self.llm, retriever.domain_mapping
        )

        self.tools = [
            Tool(
                name="domain_knowledge_retrieval",
                description="从专业领域知识库中检索相关信息",
                func=self.retrieve_domain_knowledge
            )
        ]

        self.agent = self._create_agent()

    def _create_agent(self):
        """创建Agent执行器"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个严格限定领域的知识助手，只能回答IT技术、医疗健康、人力资源三个专业领域的问题。

            处理规则：
            1. 严格判断问题是否属于三个专业领域
            2. 如果是专业领域问题：
               - 从知识库检索相关信息
               - 找到则返回优化后的答案
               - 没找到则你进行回答""
            3. 如果不是专业领域问题：
               - 直接回复"无法回答非专业领域问题"

            注意：绝对不要回答非专业领域的问题！"""),
            ("human", "{input}"),
            ("placeholder", "{agent_scratchpad}")
        ])
        return AgentExecutor(
            agent=create_tool_calling_agent(self.llm, self.tools, prompt),
            tools=self.tools,
            verbose=True
        )

    def retrieve_domain_knowledge(self, query: str) -> str:
        """知识检索工具方法"""
        domain = self.classifier.classify(query)
        if domain == "other":
            return self.retriever.out_of_domain_response

        if knowledge := self.retriever.retrieve_knowledge(query, domain):
            return knowledge
        return f"未找到{self.retriever.domain_mapping[domain]}相关专业答案"

    def answer_question(self, question: str) -> str:
        """回答用户问题"""
        try:
            response = self.agent.invoke({"input": question})["output"]
            # 确保非领域问题不会得到任何回答
            domain = self.classifier.classify(question)
            if domain == "other":
                return self.retriever.out_of_domain_response
            return response
        except Exception as e:
            print(f"回答问题出错: {str(e)}")
            return self.retriever.out_of_domain_response


# konwle = KnowledgeRetrieverFactory()
agents = AgentFactory.create_agent()

# if __name__ == "__main__":
#     # 使用工厂创建Agent
#     agent = AgentFactory.create_agent()
#
#     test_questions = [
#         "AI工具为开发者每周节省多少时间？",  # IT领域
#         "国家卫健委近期叫停哪些手术？",  # 医疗领域
#         "绩效考核周期多久合适？",  # 人力资源
#         "如何做红烧鱼？",  # 非专业领域
#         "Python怎么安装第三方库？",  # IT领域
#         "感冒了应该吃什么药？",  # 医疗领域
#         "今天天气怎么样？",  # 非专业领域
#         "请推荐一本小说"  # 非专业领域
#     ]
#
#     for q in test_questions:
#         print(f"\n问题：{q}")
#         print(f"回答：{agent.answer_question(q)}")